# coding=utf-8
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import backtrader as bt
import datetime as datetime
import pandas as pd

MIN_BARS = 30  # 交易日


# 统计回测周期内K线数量
def bar_size(file_path, from_date, to_date, date_format='%Y%m%d'):
    df = pd.read_csv(file_path)
    # 都转化为时间字符串比较，比较之后是false、true，&之后只剩下true的数据
    return len(df[(pd.to_datetime(df['date'], format=date_format) >= from_date.strftime(date_format))
                  & (pd.to_datetime(df['date'], format=date_format) <= to_date.strftime(date_format))])


# Create a Stratey
class TestStrategy(bt.Strategy):
    params = dict(
        exitbars=5,
        pfast=10,
        pslow=30
    )

    def __init__(self):
        self.sma = dict()
        self.order = None
        self.bar_executed = None
        for dat in self.datas:
            self.sma[dat._name] = bt.ind.SMA(dat.close, period=self.p.pfast)

    def log(self, txt, dt=None):
        """ Logging function for this strategy"""
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    # order 状态，为了收到订单通知，需要在用户自定义的Strategy子类中，重写notify_order方法，该方法默认为空
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # Buy/Sell order submitted/accepted to/by broker - Nothing to do
            return

        # Check if an order has been completed
        # Attention: broker could reject order if not enough cash
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log("买单成交，成交价：%.3f，成交额：%.2f，佣金：%.2f" % (
                    order.executed.price,
                    order.executed.value,
                    order.executed.comm))
            elif order.issell():
                self.log("卖单成交，成交价：%.3f，成交额：%.2f，佣金：%.2f" % (
                    order.executed.price,
                    order.executed.value,
                    order.executed.comm))
            self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log("交易失败：订单取消/资金不足/拒绝")
        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log("毛利润 %.2f, 净利润 %.2f" % (trade.pnl, trade.pnlcomm))

    def next(self):
        for data in self.datas:
            # 简单的记录收盘价作为参考
            self.log('%s 收盘 %.3f' % (data._name, data.close[0]))
            # Check if an order is pending ... if yes, we cannot send a 2nd one
            if self.order:
                return
            # 检查是否已有持仓
            if not self.position:  # 如果有仓位
                if data.close[0] > self.sma[data._name][0]:  # 今收盘 小于昨收
                    self.log("%s BUY CREATE %.3f" % (data._name, data.close[0]))
                    size = 100
                    self.buy(data=data, size=size)
            else:  # 如果没有仓位
                # Already in the market ... we might sell
                if data.close[0] < self.sma[data._name][0]:
                    self.log("%s SELL CREATE, %.2f" % (data._name, data.close[0]))
                    self.order = self.sell(data=data)


if __name__ == '__main__':
    cerebro = bt.Cerebro()

    stk_pools = ['300033', '300346']
    # 准备数据
    for stk_code in stk_pools:
        # 读取数据
        datapath = "tdx2csv/csv/sz" + stk_code + ".day.csv"
        dateformat = "%Y%m%d"  # 日期格式
        fromdate = datetime.datetime(2019, 1, 1)
        todate = datetime.datetime(2019, 12, 30)
        # 提出无效股票
        if MIN_BARS > bar_size(datapath, fromdate, todate, dateformat):  # K线数过少
            continue
        # 创建价格数据
        df_feeder = bt.feeds.GenericCSVData(  # 默认 OHLCV
            dataname=datapath,
            fromdate=fromdate,
            todate=todate,
            dtformat=dateformat,
            openinterest=-1  # -1表示该字段不是存在于CSV data
        )
        # 添加数据
        cerebro.adddata(df_feeder, name=stk_code)

    # 添加策略
    cerebro.addstrategy(TestStrategy)
    # 设置启动资金
    cerebro.broker.setcash(1000000.0)
    # 佣金 0.001 即是 0.1%
    cerebro.broker.setcommission(commission=0.001)
    # 打印初始条件
    print('Starting Portfolio Value: %.3f' % cerebro.broker.getvalue())
    # 运行回测
    cerebro.run()
    # 打印输出结果
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    # 绘图
    cerebro.plot(style='bar')
